#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File    :   bestpos_ekf_fromgnss_v2.py
@Time    :   2024/09/20 14:31:19
@Author  :   YcZ
@Version :   1.0
@Contact :   zhouyuchao@neolix.cn
@License :   (C)Copyright 2022-2025, YcZ
@Desc    :   
'''


import sys
path_to_remove = '/apollo/bazel-bin/cyber/tools/cyber_launch/cyber_launch.runfiles/apollo'
if path_to_remove in sys.path:
    sys.path.remove(path_to_remove)
    print(f'Removed {path_to_remove} from sys.path')
else:
    print(f'{path_to_remove} not found in sys.path')

from cyber.python.cyber_py3 import cyber
from modules.common_msgs.sensor_msgs.gnss_pb2 import Gnss
from modules.common_msgs.sensor_msgs.gnss_best_pose_pb2 import GnssBestPose
from filterpy.kalman import KalmanFilter as KF
import numpy as np
import pyproj


E = None

msg = GnssBestPose()

def utm2latlon(x, y, pzone=50):
    """
    convert the utm x y to lat and lon
    """
    p1 = pyproj.Proj(proj='utm', zone=pzone, ellps='WGS84')
    lon, lat = p1(x, y, inverse=True)
    return lat, lon

def latlon2utm(lon, lat, pzone=50, inverse=False):
    """
    default
        convert the lon and lat to utm x y 
    if inverse
        convert the utm x y to lon and lat
    """
    p1 = pyproj.Proj(proj='utm', zone=pzone, ellps='WGS84')
    x, y = p1(lon, lat, inverse=inverse)
    return x, y

def bestpos_callback(data):
    """
    Reader message callback.
    """
    global msg
    msg.CopyFrom(data)


def gnss_callback(data):
    """
    Reader message callback.
    """
    x = data.position.lon
    y = data.position.lat
    z = data.position.height
    x_std = data.position_std_dev.x
    y_std = data.position_std_dev.y
    z_std = data.position_std_dev.z
    x_v = data.linear_velocity.x
    y_v = data.linear_velocity.y
    z_v = data.linear_velocity.z

    x_in, y_in = latlon2utm(x, y)

    global E
    if not E:
        E = KF(dim_x=6, dim_z=6)
        E.x = np.array([x_in, y_in, z, 0, 0, 0])
        
        E.F = np.array([[1, 0, 0, 1, 0, 0],
                        [0, 1, 0, 0, 1, 0],
                        [0, 0, 1, 0, 0, 1],
                        [0, 0, 0, 1, 0, 0],
                        [0, 0, 0, 0, 1, 0],
                        [0, 0, 0, 0, 0, 1]])
        E.H = np.eye(6)
        
        E.P = np.diag([x_std ** 2, y_std ** 2, z_std ** 2, 0.01, 0.01, 0.01])

        E.R = np.diag([x_std ** 2, y_std ** 2, z_std ** 2, 0.01, 0.01, 0.01])

        E.Q = np.diag([x_std ** 2, y_std ** 2, z_std ** 2, 0.01, 0.01, 0.01])


    else:
        global bestpos_writer

        E.predict()
        z = np.array([x_in, y_in, z, x_v, y_v, z_v])
        R = np.diag([x_std ** 2, y_std ** 2, z_std ** 2, 0.01, 0.01, 0.01])
        E.update(z, R)

        x_E = E.x
        p_E = E.P

        x_out, y_out = latlon2utm(x_E[0], x_E[1], inverse=True)
        msg.longitude = x_out
        msg.latitude = y_out
        msg.height_msl = x_E[2]
        msg.latitude_std_dev = p_E[1][1] ** 0.5
        msg.longitude_std_dev = p_E[0][0] ** 0.5
        msg.height_std_dev = p_E[2][2] ** 0.5
        
        msg.header.timestamp_sec = data.header.timestamp_sec
        msg.header.sequence_num = data.header.sequence_num
        msg.measurement_time = data.measurement_time
        msg.sol_status = data.solution_status
        msg.sol_type = data.position_type

        bestpos_writer.write(msg)


def ekf_talker():

    print("=" * 80)
    bestpos_node = cyber.Node("gnss_publisher_py")
    global bestpos_writer
    bestpos_writer = bestpos_node.create_writer("/apollo/sensor/gnss/best_pose_opt", GnssBestPose, 10)

    bestpos_node.create_reader("/apollo/sensor/gnss/gnss", Gnss, gnss_callback)
    bestpos_node.create_reader("/apollo/sensor/gnss/best_pose", GnssBestPose, bestpos_callback)

    bestpos_node.spin()

def main():
    cyber.init()
    ekf_talker()

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        cyber.shutdown()
        print("\nCanceled by user.")
    finally:
        cyber.shutdown()
        print("\nExecute successfully!")


